import torch
from torch.utils.data import Dataset
from torchvision.transforms.v2 import PILToTensor,Compose
import torchvision
from torchvision import transforms, datasets
from torch.utils.data import DataLoader, RandomSampler, DistributedSampler, SequentialSampler

# 手写数字
class MNIST(Dataset):
    def __init__(self,is_train=True):
        super().__init__()
        self.ds=torchvision.datasets.MNIST('./mnist/',train=is_train,download=True)
        self.img_convert=Compose([
            PILToTensor(),
        ])
        
    def __len__(self):
        return len(self.ds)
    
    def __getitem__(self,index):
        img,label=self.ds[index]
        return self.img_convert(img)/255.0,label


def get_loader(config):
    # 数据增强
    transform_train = transforms.Compose([
        transforms.RandomResizedCrop((config.img_size, config.img_size), scale=(0.05, 1.0)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
    ])
    transform_test = transforms.Compose([
        transforms.Resize((config.img_size, config.img_size)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
    ])

    if config.dataset == "cifar10":
        trainset = datasets.CIFAR10(root="../data",
                                    train=True,
                                    download=True,
                                    transform=transform_train)
        testset = datasets.CIFAR10(root="../data",
                                   train=False,
                                   download=True,
                                   transform=transform_test)
    # RandomSampler 会对数据集进行随机抽样
    train_sampler = RandomSampler(trainset)
    test_sampler = SequentialSampler(testset)  # SequentialSampler 是按顺序逐个返回数据集中的样本
    train_loader = DataLoader(trainset,
                              sampler=train_sampler,
                              batch_size=config.train_batch_size,
                              num_workers=0,
                              pin_memory=True)
    test_loader = DataLoader(testset,
                             sampler=test_sampler,
                             batch_size=config.eval_batch_size,
                             num_workers=0,
                             pin_memory=True) if testset is not None else None

    return train_loader, test_loader



if __name__=='__main__':
    trainset = datasets.CIFAR10(root="../data",
                                    train=True,
                                    download=True)
    testset = datasets.CIFAR10(root="../data",
                                train=False,
                                download=True)
    # import matplotlib.pyplot as plt 
    
    # img_convert=Compose([
    #         PILToTensor(),
    # ])
    # trainset = datasets.CIFAR10(root="../data",
    #                                 train=True,
    #                                 download=True,)
    # x,y = trainset[3]
    # print(type(x), type(y))
    # x = img_convert(x)/255.0
    # plt.imshow(x.permute(1,2,0))
    # plt.show()
    # print(x)